"""Convert a molecule file into CPMD scripts
Usage: export_cpmd.py modelfile.mol temperature 
"""

from sys import argv
from os.path import join, split, normpath, exists, basename
from os import mkdir
from kuai.mol.io import read_molfile
from kuai.mol.element import Element
from kuai.strtool import remove_extname
from cStringIO import StringIO
from math import degrees
from shutil import copy

SYMMETRY = 1
MAXSTEP = 100000
TEMPERATURE = None
PBC = None
TIMESTEP = 2.0

ATOM_LMAX = {'H':'S', 'C':'P', 'N':'P', 'O':'P'}
ATOM_FUNCTIONS = {
    'H':('H_GIA_BLYP',   'KLEINMAN-BYLANDER'),
    'C':('C_MT_GIA_BLYP',   'KLEINMAN-BYLANDER'),
    'N':('N_MT_GIA_BLYP',   'KLEINMAN-BYLANDER'),
    'O':('O_MT_GIA_BLYP',   'KLEINMAN-BYLANDER'),
}

INPUT_OPT = """&CPMD
  OPTIMIZE WAVEFUNCTION
  PCG MINIMIZE
  TIMESTEP
  25
  COMPRESS WRITE32
  CONVERGENCE
   5.D-5 1
  EMASS
  400.0
  TASKGROUPS
  1
&END

&SYSTEM
  CHECK SYMMETRY OFF
  SYMMETRY
  %(SYMMETRY)d
  ANGSTROM
  CELL ABSOLUTE DEGREE
  %(PBC)s
  CUTOFF
  85.0
&END

%(ATOMS)s

&DFT
  NEWCODE
  FUNCTIONAL BLYP
&END

"""

INPUT_MD = """&CPMD
  MOLECULAR DYNAMICS
  RESTART WAVEFUNCTION COORDINATES    LATEST
  TEMPERATURE
  %(TEMPERATURE)f
  NOSE IONS
  %(TEMPERATURE)1f   2000.0 
  NOSE ELECTRONS
  0.02     20000.0
  MAXSTEP
  %(MAXSTEP)d
  TIMESTEP 
  %(TIMESTEP)5.2f
  EMASS
  400.0
  TRAJECTORY XYZ SAMPLE
  -10
  TASKGROUPS
  1
&END

&SYSTEM
  CHECK SYMMETRY OFF
  SYMMETRY
  %(SYMMETRY)d
  ANGSTROM
  CELL ABSOLUTE DEGREE
  %(PBC)s
  CUTOFF
  85.0
&END

%(ATOMS)s


&DFT
  NEWCODE
  FUNCTIONAL BLYP
&END
 
"""

def export(filename):
    global PBC
    mol = read_molfile(filename)
    try:
        pbc = mol.pbc
    except AttributeError:
        pbc = mol.data['PBC']
    
    pbc = pbc.a, pbc.b, pbc.c, degrees(pbc.alpha),  degrees(pbc.beta), degrees(pbc.gamma)
    PBC = '%12.6f %12.6f %12.6f %12.6f %12.6f %12.6f' % pbc
    
    workfolder = remove_extname(filename) + '.' + str(TEMPERATURE) + 'K'
    assert not exists(workfolder)
    print 'workfolder =', workfolder
    
    mkdir(workfolder)
    shname = basename(workfolder) + '.sh'
    copy(join(DATAFOLDER, 'cpmd.sh'), join(workfolder, shname))
    
    atoms = {}
    for i in mol.atoms:
        if i.symbol in atoms:
            atoms[i.symbol].append(i.coord)
        else:
            atoms[i.symbol] = [i.coord]
    
    file = StringIO()
    file.write("""&ATOMS
ISOTOPE
""")
    for k in atoms.iterkeys():
        copy(join(DATAFOLDER, ATOM_FUNCTIONS[k][0]), join(workfolder, ATOM_FUNCTIONS[k][0]))
        file.write('%12.6f\n' % Element.get(k).weight)
    file.write('\n')


    for k, v in atoms.iteritems():
        file.write('*%s %s\n' % ATOM_FUNCTIONS[k])
        file.write('   LMAX=' + ATOM_LMAX[k] + '\n')
        file.write('   ' + str(len(v)) + '\n')
        for j in v:
            file.write('%15.6f %15.6f %15.6f\n' % (j.x, j.y, j.z))
        file.write('\n\n')
    file.write('&END\n')

    global ATOMS
    ATOMS = file.getvalue()

    with open(join(workfolder, '1.inp'), 'wt') as file:
        file.write(INPUT_OPT % globals())
        
    with open(join(workfolder, '2.inp'), 'wt') as file:
        file.write(INPUT_MD % globals())

if __name__ == '__main__':
    if len(argv) < 3:
        print __doc__
    else:
        for j in argv[3:]:
            exec j
        DATAFOLDER, _ = split(argv[0])
        TEMPERATURE = int(argv[2])
        export(normpath(argv[1]))
    